{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "import metadsl.nonumpy.compat as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('__add__',\n",
       " ('array',\n",
       "  (1, 2, 3),\n",
       "  ('create_none_call',),\n",
       "  True,\n",
       "  ('create_none_call',),\n",
       "  ('create_none_call',),\n",
       "  ('create_some', 2)),\n",
       " ('arange',\n",
       "  ('create_none_call',),\n",
       "  (1, 2, 3),\n",
       "  ('create_none_call',),\n",
       "  ('create_none_call',)))"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "e = np.array((1, 2, 3), ndmin=2)\n",
    "as_tuple(e + np.arange((1, 2, 3)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'NDArray(array(NDArray((1, 2, 3)), Optional[DType](create_none_call()), Boolean(True), Optional[ArrayOrder](create_none_call()), Optional[Boolean](create_none_call()), Optional[Integer](create_some(Integer(2)))))'"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "str(e)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Tuple[Integer](shape(NDArray(array(NDArray((1, 2, 3)), Optional[DType](create_none_call()), Boolean(True), Optional[ArrayOrder](create_none_call()), Optional[Boolean](create_none_call()), Optional[Integer](create_some(Integer(2)))))))'"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "str(e.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Integer(__getitem__(Tuple[Integer](shape(NDArray(array(NDArray((1, 2, 3)), Optional[DType](create_none_call()), Boolean(True), Optional[ArrayOrder](create_none_call()), Optional[Boolean](create_none_call()), Optional[Integer](create_some(Integer(2))))))), Integer(0)))'"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "str(e.shape[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "from metadsl.expressions import Expression\n",
    "import dataclasses\n",
    "\n",
    "def as_tuple_with_type(expr):\n",
    "    return Expression.from_instance(expr).fold(\n",
    "        lambda expr: (str(expr.type), expr.value),\n",
    "        lambda call: (call.function.__name__,) + tuple(call.args),\n",
    "        lambda value: value\n",
    "    )\n",
    "\n",
    "def as_tuple(expr):\n",
    "    return Expression.from_instance(expr).fold(\n",
    "        lambda expr: expr.value,\n",
    "        lambda call: (call.function.__name__,) + tuple(call.args),\n",
    "        lambda value: value\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('array',\n",
       " (1, 2, 3),\n",
       " ('create_none_call',),\n",
       " True,\n",
       " ('create_none_call',),\n",
       " ('create_none_call',),\n",
       " ('create_some', 2))"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "as_tuple(e)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('NDArray',\n",
       " ('__add__',\n",
       "  ('NDArray',\n",
       "   ('arange',\n",
       "    ('Optional[Number]', ('create_none_call',)),\n",
       "    ('Number', 10),\n",
       "    ('Optional[Number]', ('create_none_call',)),\n",
       "    ('Optional[DType]', ('create_none_call',)))),\n",
       "  ('NDArray', 10)))"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "as_tuple_with_type(np.arange(10) + 10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "NDArray(__value__=Call(function=<function NDArray.__lt__ at 0x10e703b70>, args=(NDArray(__value__=Call(function=<function arange at 0x10e709d08>, args=(Optional(__value__=Call(function=<bound method Optional.create_none_call of <class 'metadsl.python.pure.Optional'>>, args=()), inner_type=InstanceType(type=<class 'metadsl.python.pure.Number'>, args=())), Number(__value__=10), Optional(__value__=Call(function=<bound method Optional.create_none_call of <class 'metadsl.python.pure.Optional'>>, args=()), inner_type=InstanceType(type=<class 'metadsl.python.pure.Number'>, args=())), Optional(__value__=Call(function=<bound method Optional.create_none_call of <class 'metadsl.python.pure.Optional'>>, args=()), inner_type=InstanceType(type=<class 'metadsl.nonumpy.pure.DType'>, args=()))))), NDArray(__value__=Call(function=<function NDArray.__add__ at 0x10e703510>, args=(NDArray(__value__=Call(function=<function arange at 0x10e709d08>, args=(Optional(__value__=Call(function=<function Optional.create_some at 0x10e701d90>, args=(Number(__value__=10),)), inner_type=InstanceType(type=<class 'metadsl.python.pure.Number'>, args=())), Number(__value__=100), Optional(__value__=Call(function=<bound method Optional.create_none_call of <class 'metadsl.python.pure.Optional'>>, args=()), inner_type=InstanceType(type=<class 'metadsl.python.pure.Number'>, args=())), Optional(__value__=Call(function=<function Optional.create_some at 0x10e701d90>, args=(DType(__value__='int64'),)), inner_type=InstanceType(type=<class 'metadsl.nonumpy.pure.DType'>, args=()))))), NDArray(__value__=Call(function=<function arange at 0x10e709d08>, args=(Optional(__value__=Call(function=<bound method Optional.create_none_call of <class 'metadsl.python.pure.Optional'>>, args=()), inner_type=InstanceType(type=<class 'metadsl.python.pure.Number'>, args=())), Number(__value__=10), Optional(__value__=Call(function=<bound method Optional.create_none_call of <class 'metadsl.python.pure.Optional'>>, args=()), inner_type=InstanceType(type=<class 'metadsl.python.pure.Number'>, args=())), Optional(__value__=Call(function=<bound method Optional.create_none_call of <class 'metadsl.python.pure.Optional'>>, args=()), inner_type=InstanceType(type=<class 'metadsl.nonumpy.pure.DType'>, args=())))))))))))"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = np.arange(10) < np.arange(10, 100, dtype='int64') + np.arange(10)\n",
    "\n",
    "# execute_uarray(x, backend=unumpy.backends.torch)\n",
    "x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['NUMPY_EXPERIMENTAL_ARRAY_FUNCTION'] = '1'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as real_np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('__array_function__',\n",
       " <function numpy.select(condlist, choicelist, default=0)>,\n",
       " ('from_items_call',\n",
       "  (NDArray(__value__=Call(function=<function arange at 0x10e709d08>, args=(Optional(__value__=Call(function=<bound method Optional.create_none_call of <class 'metadsl.python.pure.Optional'>>, args=()), inner_type=InstanceType(type=<class 'metadsl.python.pure.Number'>, args=())), Number(__value__=10), Optional(__value__=Call(function=<bound method Optional.create_none_call of <class 'metadsl.python.pure.Optional'>>, args=()), inner_type=InstanceType(type=<class 'metadsl.python.pure.Number'>, args=())), Optional(__value__=Call(function=<bound method Optional.create_none_call of <class 'metadsl.python.pure.Optional'>>, args=()), inner_type=InstanceType(type=<class 'metadsl.nonumpy.pure.DType'>, args=()))))),),\n",
       "  (NDArray(__value__=Call(function=<function arange at 0x10e709d08>, args=(Optional(__value__=Call(function=<bound method Optional.create_none_call of <class 'metadsl.python.pure.Optional'>>, args=()), inner_type=InstanceType(type=<class 'metadsl.python.pure.Number'>, args=())), Number(__value__=10), Optional(__value__=Call(function=<bound method Optional.create_none_call of <class 'metadsl.python.pure.Optional'>>, args=()), inner_type=InstanceType(type=<class 'metadsl.python.pure.Number'>, args=())), Optional(__value__=Call(function=<bound method Optional.create_none_call of <class 'metadsl.python.pure.Optional'>>, args=()), inner_type=InstanceType(type=<class 'metadsl.nonumpy.pure.DType'>, args=()))))),)),\n",
       " ())"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "as_tuple(real_np.select((np.arange(10),), (np.arange(10),)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('__array_ufunc__',\n",
       " <ufunc 'multiply'>,\n",
       " 'outer',\n",
       " (),\n",
       " ('arange',\n",
       "  ('create_none_call',),\n",
       "  10,\n",
       "  ('create_none_call',),\n",
       "  ('create_none_call',)),\n",
       " ('arange',\n",
       "  ('create_none_call',),\n",
       "  10,\n",
       "  ('create_none_call',),\n",
       "  ('create_none_call',)))"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "as_tuple(real_np.multiply.outer(np.arange(10), np.arange(10)))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
